rm(list = ls())

library(tidyverse)
library(lfe)
library(readxl)
library(pbapply)
library(broom)

## Helper functions

source("code/helper_functions.R")

## Data

df <- readRDS("data/mlfz_data.rds")

## Outcome list

o_list <- c(
    "TZ_u_faz_nw_bin",
    "TZ_u_bild_nw_bin",
    "TZ_u_sz_nw_bin", "TZ_u_welt_nw_bin",
    "index_nw_ratio_bin_national", "index_nw"
) %>%
    paste0(., "_mean_diff")

## Load data

df <- df %>%
  dplyr::select(county_id_2018, 
                year, 
                one_of(o_list),
                matches('covar|treat')) %>% 
  mutate(treat_categorical4 = factor(treat_categorical4,
                                     levels = c('NoChange',
                                                'DecAndInc',
                                                'Increase',
                                                'Decrease'))) %>% 
  filter(!year == 2013)

## Gen lead treatment

df <- df %>% 
  group_by(county_id_2018) %>% 
  arrange(year) %>% 
  mutate(treat_categorical4_lead1 = lead(treat_categorical4),
         treat_categorical4_lead2 = lead(treat_categorical4, 2),
         treat_categorical4_lead3 = lead(treat_categorical4, 3),
         treat_categorical4_lag1 = lag(treat_categorical4),
         treat_categorical4_lag2 = lag(treat_categorical4, 2),
         treat_categorical4_lag3 = lag(treat_categorical4, 3))

## Def covariates

covars <- c("covar_pop_total_diff", "covar_gdp_pc_diff")

## Def number of leads

lag_lead_list <- paste0(rep(c('lag', 'lead'), each = 2),
                        c(1:2, 1:2)) %>% 
  c(., "lag3")

##

res <- pblapply(o_list, function(o) {
    lapply(lag_lead_list, function(l) {
        cat(o, l, "\n")


        df_temp <- df

        ## Declare treatment

        df_temp$treatvar <- df_temp[, paste0("treat_categorical4_", l)] %>%
            pull(1)

        ## Message
        # print(o)

        ## Prep

        df_temp$outcome <- df_temp %>% pull(!!o)
        df_temp <- df_temp %>%
            filter(!is.na(outcome) & !is.na(treatvar)) %>%
            dplyr::select(
                outcome, year, county_id_2018,
                treatvar,
                one_of(covars)
            )

        ## Estimate

        m1 <- felm(outcome ~ treatvar | year | 0 | county_id_2018,
            data = df_temp
        ) %>%
            tidy_felm() %>%
            mutate(fe = "year")

        ## Year + covars

        f <- paste0(
            "outcome ~ treatvar +",
            paste0(covars, collapse = "+"),
            "| year | 0 | county_id_2018"
        ) %>%
            as.formula()

        m2 <- felm(f, data = df_temp) %>%
            tidy_felm() %>%
            mutate(fe = "year_covars")

        ## Combine

        rbind(m1, m2) %>%
            mutate(
                east = c("All"),
                outcome = !!o
            ) %>%
            mutate(period = paste0(l))
    }) %>% reduce(rbind)
}) %>%
    reduce(rbind) %>%
    filter(!str_detect(term, "Intercept|covar"))

## Save results as rds

results <- res %>% 
  mutate(term = str_remove(term,'treatvar')) %>% 
  filter(!str_detect(term, 'Intercept|covar')) %>% 
  mutate(lag = ifelse(str_detect(period, 'lag'), 'lag', 'lead')) %>% 
  mutate(period_new = as.numeric(str_extract(period, "[0-9]"))) %>% 
  mutate(period_new = ifelse(lag == 'lead', -1, 1) * period_new) %>% 
  dplyr::select(-period) %>% 
  dplyr::rename(period = period_new) %>% 
  dplyr::select(-lag)

## Covert lag / lead term to period

## Save

write_rds(
    results,
    "data/Results_MediaConsumption_Dynamic.rds"
)





